import numpy as np
import matplotlib.pyplot as plt
import gofish as gf
from astropy.io import fits
import bettermoments as bm
import pickle

path=''
line=['HDO 225.535GHz','HDO 241.558 GHz','H$2$CO (J=3$_{1,2}$-2$_{1,1}$)','C$^{17}$O J=(2-1)','CH3OH-241.806','CH3OH-241.813','CH3OH-241.852','CH3OH-241.879','CH3OH-241.887','H218O_203']
imagelist=[path+'V883_Ori_SB_HDO-225.535GHz_robust_2.0.image.fits',path+'V883_Ori_SB_HDO-241.558GHz_robust_2.0.image.fits',path+'V883_Ori_SB_H2CO-225.694GHz-fullcube_robust_2.0.image.fits',path+'V883_Ori_SB_C17O-224.711GHz_robust_2.0.image.fits', path+'V883_Ori_SB_CH3OH-241.846GHz-fullcube_robust_2.0.image.fits', path+'V883_Ori_SB_CH3OH-241.846GHz-fullcube_robust_2.0.image.fits', path+'V883_Ori_SB_CH3OH-241.846GHz-fullcube_robust_2.0.image.fits', path+'V883_Ori_SB_CH3OH-241.846GHz-fullcube_robust_2.0.image.fits', path+'V883_Ori_SB_CH3OH-241.846GHz-fullcube_robust_2.0.image.fits','V883_Ori_SB_H218O-203.GHz-0.4kms_robust_2.0.image.fits']
masklist=[path+'V883_Ori_SB_HDO-225.535GHz_robust_2.0.mask-hdo.fits',path+'V883_Ori_SB_HDO-241.558GHz_robust_2.0.mask-hdo.fits',path+'V883_Ori_SB_H2CO-225.694GHz-fullcube_robust_2.0.mask-h2co.fits',path+'V883_Ori_SB_C17O-224.711GHz_robust_2.0.mask-c17o.fits',path+'V883_Ori_SB_CH3OH-241.846GHz-fullcube_robust_2.0.mask-ch3oh-241.806.fits',path+'V883_Ori_SB_CH3OH-241.846GHz-fullcube_robust_2.0.mask-ch3oh-241.813.fits',path+'V883_Ori_SB_CH3OH-241.846GHz-fullcube_robust_2.0.mask-ch3oh-241.852.fits',path+'V883_Ori_SB_CH3OH-241.846GHz-fullcube_robust_2.0.mask-ch3oh-241.879.fits',path+'V883_Ori_SB_CH3OH-241.846GHz-fullcube_robust_2.0.mask-ch3oh-241.887.fits','V883_Ori_SB_H218O-203.GHz-0.4kms_robust_2.0.mask-h218o.fits']
firstchan=[28,31,212,72,96,155,474,691,761,31]
lastchan=[49,49,260,120,147,204,531,744,814,49]
radial_profiles={}

profile_units='K km/s'
units=['K km/s','Jy', 'Jy km/s', 'Jy/beam km/s','K']
#comment this line to run more units
units=['K km/s']

for profile_units in units:
 for i in range(len(imagelist)):
   radial_profiles[line[i]]={}
   print(imagelist[i])
   image=fits.open(imagelist[i])
   data, velax = bm.load_cube(imagelist[i])
   rms = bm.estimate_RMS(data=data, N=5)
   user_mask = bm.get_user_mask(data=data, user_mask_path=masklist[i])
   channel_mask=bm.get_channel_mask(data=data,firstchannel=firstchan[i],lastchannel=lastchan[i])
   threshold_mask = bm.get_threshold_mask(data=data,clip=2.0,smooth_threshold_mask=0.0)
   mask = bm.get_combined_mask(threshold_mask=threshold_mask,
                               channel_mask=channel_mask,
                               user_mask=user_mask,
                               combine='and')

   masked_data = data * mask
   image[0].data=masked_data
   image.writeto(imagelist[i].replace('.image.fits','.masked.image.fits'),overwrite=True)

   cube=gf.imagecube(imagelist[i].replace('.image.fits','.masked.image.fits'),FOV=3.0)

   x, y, dy = cube.radial_profile(inc=40.5, PA=32., mstar=1.285, dist=400.0, unit=profile_units)
   radial_profiles[line[i]]['x']=x.copy()
   radial_profiles[line[i]]['y']=y.copy()
   radial_profiles[line[i]]['dy']=dy.copy()
   fig, ax = plt.subplots()
   ax.errorbar(x, y, dy, fmt=' ', capsize=1.25, capthick=1.25, color='k', lw=1.0)
   ax.step(x, y, where='mid', color='k', lw=1.0)
   ax.set_xlim(x.min(), x.max())
   ax.set_xlabel('Radius (arcsec)')
   if profile_units =='K':
     ylabel='Brightness Temperature'
   elif profile_units == 'Jy':
     ylabel='Flux Density'
   elif profile_units == 'Jy km/s':
     ylabel ='Flux'
   elif profile_units == 'K km/s':
     ylabel='Intensity'
   elif profile_units == 'Jy/beam km/s':
     ylabel='Intensity'
   ax.set_ylabel(ylabel+' ('+profile_units+')')
   ax.set_title(line[i])
   plt.savefig(imagelist[i].replace('_robust_2.0.image.fits','_radial_profile_'+profile_units.replace(' ','').replace('/','')+'.png'),dpi=200.0)
   plt.savefig(imagelist[i].replace('_robust_2.0.image.fits','_radial_profile_'+profile_units.replace(' ','').replace('/','')+'.pdf'),dpi=200.0)


 pickle.dump(radial_profiles, open("radial_profiles_"+profile_units.replace(' ','').replace('/','')+".pickle", "wb"))


